#! /usr/bin/python3
# -*- coding: utf-8 -*-
"""
Created on 2020/6/2 16:06 星期二

@author: jyz
"""
import tensorflow as tf

# 1 用已有的脚本把ckpt转化为pb格式
"""
export CONFIG_FILE=gs://${YOUR_GCS_BUCKET}/data/pipeline.config
export CHECKPOINT_PATH=gs://${YOUR_GCS_BUCKET}/train/model.ckpt-2000
export OUTPUT_DIR=/tmp/tflite

python object_detection/export_tflite_ssd_graph.py \
        --pipeline_config_path=$CONFIG_FILE \
        --trained_checkpoint_prefix=$CHECKPOINT_PATH \
        --output_directory=$OUTPUT_DIR \
        --add_postprocessing_op=true
"""
# 2 用以下命令把pb转化为tflite格式。
# pb_path = ""
# converter = tf.lite.TFLiteConverter.from_frozen_graph(pb_path, ['normalized_input_image_tensor'], ['TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'], {"normalized_input_image_tensor": [1, 300, 300, 3]})
# # converter.post_training_quantize=True
# converter.optimizations = True
# converter.allow_custom_ops = True
# tflite_m = converter.convert()
#
# aa = open("converted_model.tflite", "wb")
# aa.write(tflite_m)
# aa.close()



def test():
    # 指定要使用的模型的路径  包含图结构，以及参数
    graph_def_file = '/home/docker/hello-world/research/object_detection/model_pb_fire/tflite_graph.pb'
    # 重新定义一个图
    output_graph_def = tf.GraphDef()
    with tf.gfile.GFile(graph_def_file, 'rb')as fid:
        # 将*.pb文件读入serialized_graph
        serialized_graph = fid.read()
        # 将serialized_graph的内容恢复到图中
        output_graph_def.ParseFromString(serialized_graph)
        # print(output_graph_def)
        # 将output_graph_def导入当前默认图中(加载模型)
        tf.import_graph_def(output_graph_def, name='')
    print('模型加载完成')
    # 使用默认图，此时已经加载了模型
    detection_graph = tf.get_default_graph()
    with tf.Session(graph=detection_graph)as sess:
        '''
        获取模型中的tensor
        '''
        image_tensor = detection_graph.get_tensor_by_name('input_1:0')  # pb模型输入的名字
        converter = tf.lite.TFLiteConverter.from_frozen_graph(graph_def_file,
                                                              ['normalized_input_image_tensor'],
                                                              ['TFLite_Detection_PostProcess','TFLite_Detection_PostProcess:1','TFLite_Detection_PostProcess:2','TFLite_Detection_PostProcess:3'],
                                                              {"normalized_input_image_tensor": [1, 300, 300, 3]})  # pb模型输入、输出的名字以及输入的大小
        # converter.post_training_quantize = True
        tflite_model = converter.convert()
        converter.optimizations = [tf.lite.Optimize.DEFAULT]
        open("tflite文件路径", "wb").write(tflite_model)


if __name__ == '__main__':

    # test()
    saved_model_dir = r"E:\Jie\ref_detection\model_pb_ref\saved_model"
    tf_mo = r"E:\Jie\ref_detection\model_pb_ref\converted_model.tflite"
    converter = tf.lite.TFLiteConverter.from_saved_model(saved_model_dir)
    tflite_model = converter.convert()
    open(tf_mo, "wb").write(tflite_model)


# 参考
#
# [https://blog.csdn.net/qq_26535271/article/details/84930868](https://blog.csdn.net/qq_26535271/article/details/84930868)
#
# [Tensorflow Convert pb file to TFLITE using python - Stack Overflow](https://stackoverflow.com/questions/50632152/tensorflow-convert-pb-file-to-tflite-using-python)
